import os
import sys
import subprocess
import logging
from fastapi import APIRouter, HTTPException
from typing import List, Dict, Any
from pydantic import BaseModel

router = APIRouter()
logger = logging.getLogger(__name__)

class InstallDependenciesRequest(BaseModel):
    packages: List[str]

@router.post("/install-dependencies", response_model=Dict[str, Any])
async def install_dependencies(request: InstallDependenciesRequest):
    """安装Python依赖包"""
    try:
        # 验证请求的包列表
        if not request.packages or not isinstance(request.packages, list):
            return {
                "status": "error",
                "message": "包列表不能为空"
            }
        
        # 检查包列表是否包含可接受的包
        allowed_packages = {
            # 基础包
            "pydantic", "fastapi", "uvicorn", "python-multipart", "requests", "aiofiles", 
            # LangChain相关包
            "langchain", "langchain-community", "langchain-openai", "langchain-chroma", 
            "langchain-faiss", "langchain-huggingface", 
            # 数据库和向量存储
            "chromadb", "faiss-cpu", "pypika", "duckdb", "clickhouse-connect",
            # 机器学习和向量化
            "sentence-transformers", "huggingface-hub", "transformers", "torch", "tokenizers",
            # 文档处理
            "unstructured", "pdf2image", "pytesseract", "docx2txt", "pillow",
            # 其他实用工具
            "tqdm", "numpy", "pandas", "python-dotenv"
        }
        
        for package in request.packages:
            if package.lower() not in allowed_packages:
                return {
                    "status": "error",
                    "message": f"不允许安装包: {package}"
                }
        
        # 构建pip安装命令
        cmd = [sys.executable, "-m", "pip", "install", "--upgrade"]
        cmd.extend(request.packages)
        
        # 执行命令
        logger.info(f"执行命令: {' '.join(cmd)}")
        process = subprocess.Popen(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True
        )
        
        # 获取输出
        stdout, stderr = process.communicate()
        
        # 检查命令执行结果
        if process.returncode != 0:
            logger.error(f"安装依赖失败: {stderr}")
            return {
                "status": "error",
                "message": f"安装依赖失败: {stderr}"
            }
        
        logger.info(f"安装依赖成功: {stdout}")
        return {
            "status": "success",
            "message": "依赖安装成功",
            "details": {
                "installed_packages": request.packages,
                "stdout": stdout
            }
        }
    except Exception as e:
        logger.error(f"安装依赖时发生错误: {str(e)}")
        raise HTTPException(
            status_code=500,
            detail=f"安装依赖时发生错误: {str(e)}"
        ) 